import pandas as pd
import numpy as np
from func import close_and_mcap,avg_simple_r,shares_valueadd_mkv,get_market_date,dump_to_db,dump_cash,buy_stock,drop_stock,crt1,crt2,crt3


market_date = pd.DataFrame(get_market_date('2013-01-04','2018-12-28').to_dict())
market_date = [str(i) for i in list(market_date.date)]

i = 0
change_days = []
while i<np.ceil(len(market_date)/45):
    # 每隔45个交易日进行一次主动调仓
    try:
        change_days.append(market_date[45*(i+1)])
        i += 1
    except:
        break

for date in ['2013-07-31','2013-12-13','2015-10-21','2017-04-12']:
    # 按照要求，这几个日期前后都会进行申购和赎回，因此在这几天不需要进行主动调仓
    change_days.remove(date)


purchase_redeem_info = pd.read_csv('purchase_redeem.csv')
cashflow_days = purchase_redeem_info[purchase_redeem_info['port']=='A']['date'].values

cash = 33355.0 # 建仓之后剩下的钱

market_date = pd.DataFrame(get_market_date('2013-01-04', '2018-12-28').to_dict())
market_date = [str(i) for i in list(market_date.date)]

i = 0
change_days = []
while i < np.ceil(len(market_date) / 45):
    try:
        change_days.append(market_date[45 * (i + 1)])
        i += 1
    except:
        break

for date in ['2013-07-31', '2013-12-13', '2015-10-21', '2017-04-12']:
    change_days.remove(date)

purchase_redeem_info = pd.read_csv('purchase_redeem.csv')

cashflow_days = purchase_redeem_info[purchase_redeem_info['port'] == 'A']['date'].values

cash = 33355.0  # 建仓之后剩下的钱

for i in range(len(market_date[1:])):
    today = market_date[i + 1]
    lsd = market_date[i]
    print('*' * 40, today, '*' * 40)
    lsd_info = pd.DataFrame(shares_valueadd_mkv(lsd, lsd).to_dict())  # 昨日shares、value_add和mkv

    shares = lsd_info.loc[:, ['fic_code', 'shares']]  # 昨日shares
    fic_code = tuple(set(lsd_info.fic_code))
    close_mcap = pd.DataFrame(close_and_mcap(today, fic_code).to_dict())  # 今日close和mkv

    if lsd != '2013-01-04':
        r = pd.DataFrame(avg_simple_r(fic_code, '2013-01-04', lsd).to_dict())  # 建仓开始到昨日的平均收益率
    else:
        r = pd.read_excel('stock_pool.xlsx')[['fic_code', 'r']]

    r_close_mcap = r.merge(close_mcap, on='fic_code')

    if today in change_days or today in cashflow_days:
        if today in cashflow_days:
            total_cashflow = \
            purchase_redeem_info[(purchase_redeem_info['port'] == 'A') & (purchase_redeem_info['date'] == today)][
                'cash'].values.astype('float64')
            total_cashflow = total_cashflow[0]
            print('申购赎回日,变动金额为{}'.format(float(total_cashflow)))
        else:
            total_cashflow = (np.sum(lsd_info.mkv) * 0.05)
            print('主动调仓日,变动金额为{}'.format(total_cashflow))

        if total_cashflow > 0:  # 申购
            shares_change = buy_stock(r_close_mcap, total_cashflow)
            shares = shares.merge(shares_change, on='fic_code')
            shares['shares'] = shares['shares'] + shares['shares_change']
            shares = shares.loc[:, ['fic_code', 'shares']]
            close_mcap = close_mcap.merge(shares_change, on='fic_code')
            cash_change = total_cashflow - np.sum(close_mcap['close'] * close_mcap['shares_change'])
            close_mcap = close_mcap.loc[:, ['fic_code', 'close', 'mktcap']]
            shares_close_mcap = close_mcap.merge(shares, on='fic_code')
            cash += cash_change
            print('申购成功！！！！！！！！！！！！！！')

        else:  # 赎回 total_cashflow<0
            total_cashflow = -total_cashflow
            print('赎回金额为{}'.format(total_cashflow))
            if total_cashflow > cash:  # 如果剩余现金不够赎回
                total_cashflow -= cash
                cash = 0
                print('开始drop')
                shares_change = drop_stock(r_close_mcap, shares, total_cashflow, 0, None)
                print(shares_change)
                shares = shares.merge(shares_change, on='fic_code')
                shares['shares'] = shares['shares'] - shares['shares_change']
                shares = shares.loc[:, ['fic_code', 'shares']]
                close_mcap = close_mcap.merge(shares_change, on='fic_code')
                cash_change = np.sum(close_mcap['close'] * close_mcap['shares_change']) - total_cashflow
                close_mcap = close_mcap.loc[:, ['fic_code', 'close', 'mktcap']]
                shares_close_mcap = close_mcap.merge(shares, on='fic_code')
                cash += cash_change
            else:
                cash -= total_cashflow  # 如果现金够赎回
            print('赎回成功!!!!!!!!!!')

    shares_close_mcap = close_mcap.merge(shares, on='fic_code')

    crt1_result = crt1(shares_close_mcap, cash)
    ctr = len(crt2(shares_close_mcap)) + len(crt3(shares_close_mcap)) + int(crt1_result)
    while ctr > 0:
        stock_need_drop = set((crt2(shares_close_mcap))) | set(crt3(shares_close_mcap))
        print('不满足条件，需要减仓的股票是:{}'.format(stock_need_drop))
        shares_change = drop_stock(r_close_mcap, shares_close_mcap[['fic_code', 'shares']], 0, 1, stock_need_drop)

        shares_close_mcap = shares_close_mcap.merge(shares_change, on='fic_code')
        shares_close_mcap['shares'] = shares_close_mcap['shares'] - shares_close_mcap['shares_change']
        cash += (shares_close_mcap['close'] * shares_close_mcap['shares_change']).sum()
        shares_close_mcap = shares_close_mcap.iloc[:, :-1]  # 把shares_change这列删除
        print('减仓后，现金为{}'.format(cash))

        # 现金增加，购买股票
        shares_change = buy_stock(r_close_mcap, cash)
        shares_close_mcap = shares_close_mcap.merge(shares_change, on='fic_code')
        shares_close_mcap['shares'] = shares_close_mcap['shares'] + shares_close_mcap['shares_change']
        cash -= (shares_close_mcap['close'] * shares_close_mcap['shares_change']).sum()
        print('加仓共花费:{}'.format((shares_close_mcap['close'] * shares_close_mcap['shares_change']).sum()))
        shares_close_mcap = shares_close_mcap.iloc[:, :-1]  # 把shares_change这列删除

        ctr = int(len(crt2(shares_close_mcap)) > 0) + int(len(crt3(shares_close_mcap)) > 0) + int(crt1_result)
    else:
        print('符合要求')
        port_mkv = np.sum(shares_close_mcap['shares'] * shares_close_mcap['close'])

    print('今天是{},port的mkv:{},cash:{},port_mkv/cash:{}'.format(today, port_mkv, cash, port_mkv / cash))
    lsd_mkv = lsd_info.loc[:, ['fic_code', 'mkv']]
    lsd_shares = lsd_info.loc[:, ['fic_code', 'shares']]
    lsd_shares.columns = ['fic_code', 'lsd_shares']
    shares_gap = shares_close_mcap.merge(lsd_shares, on='fic_code')
    shares_gap['shares_gap'] = shares_close_mcap['shares'] - shares_gap['lsd_shares']
    dump_info = shares_close_mcap.merge(lsd_mkv, on='fic_code')
    dump_info['security_mkt_value'] = dump_info['shares'] * dump_info['close']
    dump_info['cash_flow'] = dump_info.close * shares_gap['shares_gap']
    dump_info['value_add'] = dump_info['security_mkt_value'] - dump_info['mkv'] - dump_info['cash_flow']
    dump_info['weight'] = dump_info.security_mkt_value / np.sum(dump_info.security_mkt_value)
    dump_info['portfolio_id'] = 110001
    dump_info['date'] = today
    dump_info = dump_info.to_dict('records')
    dump_to_db(dump_info)
    dump_cash(today, cash)